{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "35813d22-9f27-4d87-8623-53397599409a",
   "metadata": {},
   "source": [
    "### 原始生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fc31193-4f3d-49f9-adea-3c74f52de9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "model_name = \"llama-2-7b-hf\" # 用你下载的模型的文件夹位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "\n",
    "logits = model.forward(input_ids)\n",
    "print(\"Logits Shape:\", logits.logits.shape)\n",
    "print(f\"logits:{logits.logits}\")\n",
    "\n",
    "next_token = torch.argmax(logits.logits, dim=-1).reshape(-1)[1]\n",
    "print(f\"next_token:{next_token}\")\n",
    "\n",
    "next_word = tokenizer.decode(next_token)\n",
    "print(f\"next_word:{next_word}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3f4e0e0-3057-49bc-aedd-4704b021408e",
   "metadata": {},
   "source": [
    "### temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1a96f09f-0799-46af-9627-ef03b9844039",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "probs:tensor([[0.2559, 0.5154, 0.0571, 0.1716]])\n",
      "probs_low:tensor([[0.1800, 0.7301, 0.0090, 0.0809]])\n",
      "probs_high:tensor([[0.2695, 0.3825, 0.1273, 0.2207]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "logits = torch.tensor([[0.5, 1.2, -1.0, 0.1]])\n",
    "# 无temperature\n",
    "probs = torch.softmax(logits, dim=-1)\n",
    "# temperature low 0.5\n",
    "probs_low = torch.softmax(logits / 0.5, dim=-1)\n",
    "# temperature high 2\n",
    "probs_high = torch.softmax(logits / 2, dim=-1)\n",
    "\n",
    "print(f\"probs:{probs}\")\n",
    "print(f\"probs_low:{probs_low}\")\n",
    "print(f\"probs_high:{probs_high}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2d69b9-0209-405d-94a6-800fa4622459",
   "metadata": {},
   "source": [
    "### top_p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e77008-380e-4a23-a609-c2a58b639873",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "# 样例：probs: tensor([[0.2559, 0.5154, 0.0571, 0.1716]])\n",
    "probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])\n",
    "# 第一步进行排序\n",
    "probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)\n",
    "# 结果\n",
    "probs_sort: tensor([[0.5154, 0.2559, 0.1716, 0.0571]])\n",
    "probs_idx: tensor([[1, 0, 3, 2]])\n",
    "\n",
    "# 第二步概率的累积和\n",
    "probs_sum = torch.cumsum(probs_sort, dim=-1)\n",
    "# 结果\n",
    "probs_sum: tensor([[0.5154, 0.7713, 0.9429, 1.0000]])\n",
    "\n",
    "# 第三步找到第一个大于阈值p的位置，假设p=0.9，并将后面的概率值置为0：\n",
    "mask = probs_sum - probs_sort > p\n",
    "probs_sort[mask] = 0.0\n",
    "# 结果\n",
    "probs_sort: tensor([[0.5154, 0.2559, 0.1716, 0.0000]])\n",
    "\n",
    "# 第四步复原原序列\n",
    "new_probs = probs_sort.scatter(1, probs_idx, probs_sort)\n",
    "# 结果\n",
    "new_probs: tensor([[0.2559, 0.5154, 0.0000, 0.1716]])\n",
    "\n",
    "# 注：在真实实现中一般会把舍弃的概率置为-inf，即\n",
    "zero_indices = (new_probs == 0)\n",
    "new_probs[zero_indices] = float('-inf')\n",
    "# 结果\n",
    "new_probs: tensor([[0.2559, 0.5154, -inf, 0.1716]])\n",
    "# 完整代码\n",
    "def sample_top_p(probs, p):\n",
    "    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)\n",
    "    probs_sum = torch.cumsum(probs_sort, dim=-1)\n",
    "    mask = probs_sum - probs_sort > p\n",
    "    probs_sort[mask] = 0.0\n",
    "    new_probs = probs_sort.scatter(1, probs_idx, probs_sort)\n",
    "    zero_indices = (new_probs == 0)\n",
    "    new_probs[zero_indices] = float('-inf')\n",
    "    return new_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb3d974-75e3-43df-915a-eb22771f86b3",
   "metadata": {},
   "source": [
    "### top_k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c575d66f-059d-49ea-81e8-5ba948521254",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new_probs: tensor([[0.2559, 0.5154,   -inf,   -inf]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "filter_value = -float(\"Inf\")\n",
    "top_k = 2\n",
    "probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])\n",
    "indices_to_remove = probs < torch.topk(probs, top_k)[0][..., -1, None]\n",
    "new_probs = probs.masked_fill(indices_to_remove, filter_value)\n",
    "print(\"new_probs:\", new_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88402e6a-2909-4628-9d4c-4111053f38b9",
   "metadata": {},
   "source": [
    "### repetition_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0b4101f-a916-4068-b1cf-f2c328429b07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "原始概率分布： [0.3 0.1 0.3 0.1 0.2]\n",
      "调整后的概率分布： [0.33333333 0.11111111 0.26666667 0.11111111 0.17777778]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "def apply_repetition_penalty(probs, repetition_penalty, prev_tokens):\n",
    "    adjusted_probs = np.copy(probs)\n",
    "    for token in set(prev_tokens):\n",
    "        adjusted_probs[token] *= repetition_penalty\n",
    "    adjusted_probs /= np.sum(adjusted_probs)  \n",
    "    return adjusted_probs\n",
    "# 示例概率分布，索引对应不同词语\n",
    "original_probs = np.array([0.3, 0.1, 0.3, 0.1, 0.2])\n",
    "# 示例先前生成的词语\n",
    "previous_tokens = [2, 4, 2]\n",
    "# 重复惩罚系数\n",
    "repetition_penalty = 0.8\n",
    "# 应用重复惩罚，得到调整后的概率分布\n",
    "adjusted_probs = apply_repetition_penalty(original_probs, repetition_penalty, previous_tokens)\n",
    "\n",
    "print(\"原始概率分布：\", original_probs)\n",
    "print(\"调整后的概率分布：\", adjusted_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "152f8cdb-97f9-4a66-9206-c82f9306045b",
   "metadata": {},
   "source": [
    "### do_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d1b6b1a8-ccf8-49d6-90b2-6a8b79ce0ad5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "next_token: tensor([[1]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "probs = torch.tensor([[0.2559, 0.5154, 0.0571, 0.1716]])\n",
    "next_token = torch.multinomial(probs, num_samples=1)\n",
    "print(\"next_token:\", next_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "714e70d5-bbf2-46c8-97c4-b6a0428e065a",
   "metadata": {},
   "source": [
    "### num_beams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cc0c0b32-d922-4ee5-98a8-dedcf0ab49fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence 1: <start> I like\n",
      "Sentence 2: <start> I apple\n",
      "Sentence 3: <start> I peach\n"
     ]
    }
   ],
   "source": [
    "class BeamSearchNode:\n",
    "    def __init__(self, sequence, score):\n",
    "        self.sequence = sequence  # 生成的序列\n",
    "        self.score = score  # 分数（概率）\n",
    "        \n",
    "# 示例：下一个token的概率函数，简单使用固定概率\n",
    "def simple_next_word_probs(sequence):\n",
    "    if sequence[-1] == \"<end>\":\n",
    "        return {}\n",
    "    return {\"apple\": 0.3, \"like\": 0.35, \"peach\": 0.2, \"banana\": 0.15}\n",
    "\n",
    "\n",
    "def beam_search(initial_sequence, next_word_probs_func, num_beams, max_sequence_length):\n",
    "    # 初始化初始节点，且分数为1\n",
    "    initial_node = BeamSearchNode(sequence=initial_sequence, score=1.0)\n",
    "    candidates = [initial_node]\n",
    "\n",
    "    final_candidates = []  # 最终的候选序列\n",
    "    # 只要候选节点列表不为空，且 final_candidates 中的候选节点数量还没有达到指定的束宽度，就继续进行搜索\n",
    "    while candidates and len(final_candidates) < num_beams:\n",
    "        # 候选节点排序\n",
    "        candidates.sort(key=lambda x: -x.score)\n",
    "        current_node = candidates.pop(0)\n",
    "        # 当节点序列末尾生成结束符号（如\"<end>\"），或者当生成的序列长度达到最大限制时终止节点的扩展\n",
    "        if current_node.sequence[-1] == \"<end>\" or len(current_node.sequence) >= max_sequence_length:\n",
    "            final_candidates.append(current_node)\n",
    "        else:\n",
    "            # 获取下一个token的概率，我们的例子返回的是固定的概率\n",
    "            next_words_probs = next_word_probs_func(current_node.sequence) \n",
    "            # 生成新的候选序列，并计算分数           \n",
    "            for next_word, next_word_prob in next_words_probs.items():\n",
    "                new_sequence = current_node.sequence + [next_word]\n",
    "                new_score = current_node.score * next_word_prob\n",
    "                new_node = BeamSearchNode(sequence=new_sequence, score=new_score)\n",
    "                candidates.append(new_node)\n",
    "\n",
    "    return [candidate.sequence for candidate in final_candidates]\n",
    "\n",
    "# 开始使用：\n",
    "\n",
    "initial_sequence = [\"<start>\", \"I\"]\n",
    "num_beams = 3\n",
    "max_sequence_length = 3\n",
    "result = beam_search(initial_sequence, simple_next_word_probs, num_beams, max_sequence_length)\n",
    "\n",
    "for idx, sequence in enumerate(result):\n",
    "    print(f\"Sentence {idx + 1}: {' '.join(sequence)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fa70db8-d21e-4e7c-b0d0-7e4f4ad17e08",
   "metadata": {},
   "source": [
    "### constrained beam-search decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "746241eb-5f34-47bd-af4a-ff1a29811db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "import torch\n",
    "model_name = \"llama-2-7b-hf\" # 你模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "# generate实现\n",
    "generation_output = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    num_beams = 3,\n",
    "    num_return_sequences=3,\n",
    "    return_dict_in_generate=True,\n",
    "    max_new_tokens=3,\n",
    ")\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deb0a45c-2255-4254-b8e5-a0748b5f71ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "import torch\n",
    "model_name = \"llama-2-7b-hf\" # 你模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "force_words = [\"my\"]\n",
    "force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids\n",
    "\n",
    "generation_output = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    force_words_ids = force_words_ids,\n",
    "    num_beams = 3,\n",
    "    num_return_sequences=3,\n",
    "    return_dict_in_generate=True,\n",
    "    max_new_tokens=3,\n",
    ")\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0589bee-997d-449b-9db3-a6b21b2c530c",
   "metadata": {},
   "source": [
    "### contrastive search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "831debf4-b0d8-4c6f-b7b0-710d92286d00",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "import torch\n",
    "model_name = \"llama-2-7b-hf\" # 你模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "\n",
    "generation_output = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    penalty_alpha = 0.5,\n",
    "    top_k = 30,\n",
    "    return_dict_in_generate=True,\n",
    "    max_new_tokens=3,\n",
    ")\n",
    "\n",
    "# 直接使用其函数\n",
    "# generation_output = model.contrastive_search(\n",
    "#     input_ids=input_ids,\n",
    "#     penalty_alpha = 0.5,\n",
    "#     top_k = 30,\n",
    "#     return_dict_in_generate=True,\n",
    "#     max_new_tokens=3,\n",
    "# )\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749fa43c-4712-4e35-92d5-c29e8cc5f940",
   "metadata": {},
   "source": [
    "### greedy decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb414f73-7d16-4262-99b9-ce104c2add78",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "import torch\n",
    "model_name = \"llama-2-7b-hf\" # 你模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "\n",
    "generation_output = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    num_beams = 1,\n",
    "    do_sample = False,\n",
    "    return_dict_in_generate=True,\n",
    "    max_new_tokens=3,\n",
    ")\n",
    "# 直接指定使用其函数\n",
    "# generation_output = model.greedy_search(\n",
    "#     input_ids=input_ids,\n",
    "#     num_beams = 1,\n",
    "#     do_sample = False,\n",
    "#     return_dict_in_generate=True,\n",
    "#     max_length = 7\n",
    "# )\n",
    "\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bd85cb9-0cb2-41d8-ad8e-a1c24fa3e2bb",
   "metadata": {},
   "source": [
    "### multinomial sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2870cf87-304a-471a-a35a-c63fe714edf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "from transformers import (\n",
    "    LogitsProcessorList,\n",
    "    TopKLogitsWarper,\n",
    "    TopPLogitsWarper,\n",
    "    TemperatureLogitsWarper,\n",
    "    )\n",
    "\n",
    "import torch\n",
    "model_name = \"llama-2-7b-hf\" # 你模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "\n",
    "generation_output = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    num_beams = 1,\n",
    "    do_sample = True,\n",
    "    temperature = 1.2,\n",
    "    top_k = 100,\n",
    "    top_p = 0.6,\n",
    "    return_dict_in_generate=True,\n",
    "    max_length=7,\n",
    ")\n",
    "\n",
    "\n",
    "# sample实现\n",
    "# logits_warper = LogitsProcessorList(\n",
    "#     [\n",
    "#     TopKLogitsWarper(100),\n",
    "#     TemperatureLogitsWarper(1.2),\n",
    "#     TopPLogitsWarper(0.6)\n",
    "#     ]\n",
    "# )\n",
    "# generation_output = model.sample(\n",
    "#     input_ids=input_ids,\n",
    "#     logits_warper=logits_warper,\n",
    "#     return_dict_in_generate=True,\n",
    "#     max_length=7,\n",
    "# )\n",
    "\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf86750-0f17-49a0-a9fd-0e94b355e4d0",
   "metadata": {},
   "source": [
    "### assisted decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2dad77-d546-4460-9536-1e1c1232d695",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "import torch\n",
    "model_name = \"llama-2-13b-hf\" # 你自己模型的位置\n",
    "assistant_model_name = \"llama-2-7b-hf\" # 你自己模型的位置\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\")\n",
    "assistant_model = AutoModelForCausalLM.from_pretrained(assistant_model_name, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "text = \"say hello to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "print(f\"inputs:{inputs}\")\n",
    "input_ids = inputs[\"input_ids\"].to(\"cuda\")\n",
    "\n",
    "\n",
    "generation_output = model.generate(\n",
    "    assistant_model=assistant_model,\n",
    "    input_ids=input_ids,\n",
    "    num_beams = 1,\n",
    "    do_sample = False,\n",
    "    return_dict_in_generate=True,\n",
    "    max_length=7,\n",
    ")\n",
    "\n",
    "\n",
    "print(\"query:\", text)\n",
    "for i, output_sequence in enumerate(generation_output.sequences):\n",
    "    output_text = tokenizer.decode(output_sequence, skip_special_tokens=True)\n",
    "    print(f\"Generated sequence {i+1}: {output_text}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
